'''
第四章重要的分析部分

用于分析图像中，并且计算单个图层的贡献度，
这里会创建不同的图像结果。
'''
from AnalysConv.outputWeight import SaveNumpy_Weight
from AnalysConv.AnalysConv import splitConvLayer,CalConvoluationLayer,CaluationConvlayers,SaveCaluationResult
from AnalysConv.MergeExcel import outConv1ToXls_matrix,outConv1ToXls_Single
import numpy as np
import os
import sys

from torch.nn.modules import conv
from matplotlib import pyplot as plt
import math
import torch
import torch.nn as torch_nn
import torch.nn.functional as F
from models.get_models import get_models
from utils.utils import load_ckpt,getedge,getedge_sig


# ----方法组-----
def LoadModel(modelName,modelcpk_path,num_cls,input_channels):
    '''
    加载模型
    args:
        modelName 需要参与计算的模型的名称
        modelcpk_path 训练好模型的保存路径
    return:
          model : 加载好对应权重的模型
    '''
    model=get_models(modelName,num_class=num_cls,channels_num=input_channels)
    epoch,model,best_iou=load_ckpt(model,None,modelcpk_path,'cuda')
    return epoch,model,best_iou

def preProcess(image,mean_arr,std_arr):
    '''
    图像 zeros 归一化
    '''
    image=(image-mean_arr)/std_arr # z-score 零-均值归一化
    return image 

def getImage(imginfo):
    '''
    返回图像预处理的结果
    '''
    img=np.load(imginfo['img_path'])
    label=np.load(imginfo['label_path'])
    seg=np.load(imginfo['seg_path'])
    mean_arr=np.array(imginfo['mean'])
    std_arr=np.array(imginfo['std'])    
    img_process=preProcess(img,mean_arr,std_arr) 
    img_dict={'ori_img':img,'pre_img':img_process,'label':label,'seg':seg}
    return img_dict


def mainprocess(img_dict,conv_modules,root_dir,img_name="test"):
    '''
    主要的计算流程
    '''
    img_input=img_dict['pre_img'] # 预处理
    convlayers=splitConvLayer(conv_modules) # 分割成两层
    # 保存权重文件，生成xlxs 文件，相应的计算层情况
    SaveNumpy_Weight(os.path.join(root_dir,'weights'),convlayers)
    outConv1ToXls_matrix(os.path.join(root_dir,'convlayer_matrix.xlsx'),convlayers)
    outConv1ToXls_Single(os.path.join(root_dir,'convlayer_single.xlsx'),convlayers)
    img_input=np.transpose(img_input,(2,0,1))
    
    cal_result=CaluationConvlayers(img_input,convlayers) # 计算实际结果
    calresult_dir=os.path.join(root_dir,img_name)
    SaveCaluationResult(calresult_dir,cal_result) # 保存计算结果
    # 保存贡献度到不同的文件中的

# ----主要调用方法-----
if __name__=='__main__':
    # 加载模型
    modelinfo={
        'modelname':'unet',
        'modelcpk':'/media/gis/databackup/ayc/model_hastrained_log/RGB234_unet_result/checkpoint/Bestest_mIou_ckpt_unet_dataName_RGB234.pth',
        'num_cls':2,
        'input_channels':3
    }

    imginfo={
        'img_path':'/media/gis/databackup/ayc/modellist/dataset/nanchang/E512/RGB234/test/img/2037.npy',
        'label_path':'/media/gis/databackup/ayc/modellist/dataset/nanchang/E512/RGB234/test/label/2037.npy',
        'seg_path':'/media/gis/databackup/ayc/modellist/dataset/nanchang/E512/RGB234/test/seg/2037.npy',
        'mean':[103.47446405573739, 89.24539776311728, 118.10080357709782],
        'std':[40.96571881718136, 36.232309241791654, 46.13046923546727]
    }
    # 结果的输出根目录
    result_rootdir="/media/gis/databackup/ayc/modellist/AnalysConv/chap4"
    # modelinfo params
    modelName=modelinfo['modelname']
    modelcpk_path=modelinfo['modelcpk']
    num_cls=modelinfo['num_cls']
    input_channels=modelinfo['input_channels']
    epoch,model_unet,best_iou=LoadModel(modelName,modelcpk_path,num_cls,input_channels)
    # 抽出其中的最为关键的两层模型
    '''
    model_unet inc.conv 部分模块结构
    0:Conv2d
    1:BatchNorm2d
    2:ReLu
    3:Conv
    4:BatchNorm2d
    5:ReLU
    这里做个简单的论证，其中的BatchNorm2d 与ReLU 没有改变 激活值的分布情况，不会改变计算结果的大小关系。
    这是因为这两种函数都是单调递增关系，因此只是对于由卷积计算的结果进行进一步的缩放，
    但是不会修改数据本身所具有的大小关系与拓扑顺序，只是调整了数据本身的分布情况。
    而单个特征波段对于最终激活值的影响也并不会随着正则与激活两种计算方法而产生影响，即正则运算与激活运算与卷积运算的结果之间是独立且互不影响的。
    '''
    # 结构转换
    conv_modules={'conv_1':model_unet.inc.conv.conv._modules['0'],
                    'batchNorm2d_1':model_unet.inc.conv.conv._modules['1'],
                    'ReLU_1':model_unet.inc.conv.conv._modules['2'],
                    'conv_2':model_unet.inc.conv.conv._modules['3'],
                    'batchNorm2d_2':model_unet.inc.conv.conv._modules['4'],
                    'ReLU_2':model_unet.inc.conv.conv._modules['5']}
    
    # 处理图像预处理
    img_dict=getImage(imginfo)
    # 计算图像结果
    mainprocess(img_dict,conv_modules,result_rootdir)

    